# Copyright (c) Alibaba, Inc. and its affiliates.
# Part of the implementation is borrowed from huggingface/transformers.
import inspect
import os
import shutil
import time
from contextlib import contextmanager
from copy import copy
from functools import partial
from types import MethodType
from typing import Callable, Dict, List, Optional, Tuple, Union

import safetensors
import torch
import torch.distributed as dist
import torch.nn as nn
import transformers
from datasets import Dataset as HfDataset
from modelscope import check_local_model_is_latest
from packaging import version
from peft import PeftModel
from torch.nn import Module
from torch.utils.data import DataLoader
from transformers import PreTrainedModel
from transformers.data.data_collator import DataCollator
from transformers.integrations import is_deepspeed_zero3_enabled
from transformers.modeling_utils import unwrap_model
from transformers.trainer import TrainerCallback
from transformers.trainer_utils import EvalPrediction, IntervalStrategy
from transformers.utils import is_torch_npu_available

from swift.hub import get_hub
from swift.llm import BatchSamplerShard, DataLoaderDispatcher, DataLoaderShard, Template
from swift.plugin import MeanMetric, compute_acc, extra_tuners
from swift.tuners import SwiftModel
from swift.utils import get_logger, is_mp_ddp, use_torchacc
from swift.utils.torchacc_utils import ta_trim_graph
from ..utils.torch_utils import get_device_count
from .arguments import TrainingArguments
from .utils import can_return_loss, find_labels, get_function, is_instance_of_ms_model

try:
    from trl import AutoModelForCausalLMWithValueHead
except (ImportError, RuntimeError):
    AutoModelForCausalLMWithValueHead = None

logger = get_logger()


class SwiftMixin:
    """
    SwiftMixin 是一个针对训练器（Trainer）设计的混入类（Mixin），
    它扩展了模型训练、保存、优化器管理、多模态支持、分布式训练兼容性、日志记录等功能，特别针对 Swift 框架、PEFT（参数高效微调）以及相关模型架构进行了定制优化。
    其设计目的是使得训练流程更加稳定、灵活、支持多模态和分布式训练，且方便保存模型和上传到模型库（Hub）。
    """

    def __init__(self,
                 model: Union[PreTrainedModel, Module] = None,
                 args: TrainingArguments = None,
                 data_collator: Optional[DataCollator] = None,
                 train_dataset: Optional[HfDataset] = None,
                 eval_dataset: Optional[Union[HfDataset, Dict[str, HfDataset]]] = None,
                 template: Optional[Template] = None,
                 model_init: Optional[Callable[[], PreTrainedModel]] = None,
                 compute_loss_func: Optional[Callable] = None,
                 compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
                 callbacks: Optional[List[TrainerCallback]] = None,
                 optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
                 preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,
                 **kwargs) -> None:
        """
        初始化 SwiftMixin。

        - 检查模型版本是否最新（如果 args.check_model 且 model 有 model_dir 属性）
        - 禁用 eval_dataset 时关闭评估策略
        - 初始化自定义指标容器 _custom_metrics
        - 存储模板和模型元信息，初始化 hub 交互对象
        - 调用父类初始化，传入必要参数和模板的 tokenizer
        - 处理前向函数是否被重写的判断，获取标签名和是否可以返回 loss
        - 如果模板启用 sequence_parallel_size > 1，加载 sequence_parallel 支持

        model: 训练使用的模型（支持 HuggingFace 的 PreTrainedModel 或自定义模块 torch.nn.Module）
        args: 训练参数（TrainingArguments）
        data_collator: 批量数据处理器（可选，用于自定义 batch 构造方式）
        train_dataset: 训练集
        eval_dataset: 验证集（支持字典格式用于多任务）
        template: Swift 框架中的模板类，用于注入 tokenizer 和配置
        model_init: 可选的模型初始化函数（常用于重新初始化）
        compute_loss_func: 自定义损失函数
        compute_metrics: 自定义评估指标函数
        callbacks: 训练回调列表
        optimizers: 优化器和学习率调度器（tuple）
        preprocess_logits_for_metrics: 用于评估时对 logits 的预处理函数
        **kwargs: 其他传递给父类的参数
        """

        # 如果训练参数开启了模型检查，且模型含有 model_dir 属性（一般用于 Swift 本地模型管理）
        if args.check_model and hasattr(model, 'model_dir'):
            from swift.utils.logger import ms_logger_ignore_error
            with ms_logger_ignore_error():
                check_local_model_is_latest(
                    model.model_dir, user_agent={
                        'invoked_by': 'local_trainer',
                        'third_party': 'swift',
                    })

        # 如果未提供验证集，手动将评估策略设为禁用（防止报错）
        if eval_dataset is None and args:
            args.evaluation_strategy = IntervalStrategy.NO
            args.eval_strategy = IntervalStrategy.NO

        # 初始化一个字典，用于在训练过程中记录自定义指标
        self._custom_metrics = {}
        # 存储传入的模板类（通常含有 tokenizer、数据预处理等逻辑）
        self.template = template
        # 初始化显存占用记录（用于后续统计最大显存）
        self.max_memory = 0
        # 获取模型上传/下载的中介对象 hub（用于和模型库交互）
        self.hub = get_hub()

        # 获取模型元信息（如模型类型、作者、版本等，需 model 具备 model_meta 属性）
        self.model_meta = model.model_meta
        # 在 patch_hub 上下文中执行操作，通常用于 monkey-patch 特定逻辑（如上传格式）
        with self.hub.patch_hub():
            super().__init__(
                model=model,
                args=args,
                data_collator=data_collator,
                train_dataset=train_dataset,
                eval_dataset=eval_dataset,
                tokenizer=template.tokenizer,
                model_init=model_init,
                compute_metrics=compute_metrics,
                callbacks=callbacks,
                optimizers=optimizers,
                preprocess_logits_for_metrics=preprocess_logits_for_metrics,
                **kwargs)

        # 将自定义损失函数保存在类中，用于后续覆盖默认的损失计算逻辑
        self.compute_loss_func = compute_loss_func

        # 判断模型的 forward 函数是否被重写。
        # 目的：有些模型可能重写了 forward，需要特殊处理标签名或 loss。
        if get_function(model.__class__.forward) is not get_function(model.forward):
            # 自动查找模型中用于监督训练的标签字段名；
            self.label_names = find_labels(model)
            # 判断模型是否可以直接返回 loss（即 forward 的输出是否包含 loss 字段）
            self.can_return_loss = can_return_loss(model)

        # 如果未能成功自动识别标签名，则默认为 ['labels']
        self.label_names = self.label_names or ['labels']

        # 记录训练开始时间（用于日志记录和性能分析）
        self.start_time = time.time()

        # 如果模板启用了序列并行（通常用于大型模型多卡训练），则引入并初始化并行逻辑。
        if self.template.sequence_parallel_size > 1:
            from swift.trainers.sequence_parallel import sequence_parallel
            # sequence_parallel.prepare_trainer(self) 会修改模型结构、分布式策略等。
            sequence_parallel.prepare_trainer(self)

    def _save_initial_model(self, output_dir):
        """
        保存模型的初始权重，专用于 LoRA 技术（如 pissa、olora、lora-ga）。

        用途：
        - 某些 LoRA 扩展技术（如 pissa）在训练时需要保存初始权重；
        - 该方法会在训练开始前调用一次，用于后续转换或融合操作；
        - 仅当模型是 PeftModel 且启用了 init_lora_weights 字段时才触发保存。
        """

        # 解包模型，移除可能的封装（如 DDP 或 fp16 包裹）以获取底层 PeftModel
        model = unwrap_model(self.model)

        # 判断模型是否是 PEFT 结构（如 LoRA 包裹的模型）
        if isinstance(model, PeftModel):
            # 获取默认的 LoRA 配置（可能是多个适配器之一）
            config = model.peft_config.get('default')

            # 获取 init_lora_weights 字段（指示是否需要保存初始 LoRA 权重）
            init_lora_weights = getattr(config, 'init_lora_weights', None)

            # 如果 init_lora_weights 是字符串，且包含特定关键字（pissa、olora、lora-ga）
            if (isinstance(init_lora_weights, str)
                    and any(s in init_lora_weights for s in ('pissa', 'olora', 'lora-ga'))):
                # 设置为 True，表示正在保存初始 LoRA 权重（临时更改该字段）
                config.init_lora_weights = True

                # 保存当前 PeftModel 的初始权重到指定路径下的 'initial_model' 子目录
                model.save_pretrained(os.path.join(output_dir, 'initial_model'))

                # 恢复 config 中原始的 init_lora_weights 配置，避免影响后续流程
                config.init_lora_weights = init_lora_weights

    def _save_converted_model(self, output_dir):
        """
        保存转换后的模型权重（如 LoRA-GA、pissa、olora 等方法转换后的权重）。

        功能说明：
        - 根据配置参数 `init_lora_weights` 判断是否需要进行特殊的权重转换保存；
        - 如果使用了 LoRA-GA，需要引入其专用上下文管理器进行转换；
        - 将转换后的权重保存到 output_dir/converted 文件夹；
        - 支持三种 LoRA 扩展方法：LoRA-GA、pissa、olora。
        """

        # 解包模型，移除可能存在的 DDP、FP16、混合精度等封装，得到底层原始模型
        model = unwrap_model(self.model)

        # 判断模型是否为 PeftModel（即使用 PEFT 框架封装的模型，如 LoRA、AdaLoRA 等）
        if isinstance(model, PeftModel):

            # 获取默认 adapter（通常 key 为 'default'）对应的配置对象
            config = model.peft_config.get('default')

            # 获取是否启用了 init_lora_weights，通常是一个标识字符串（如 'pissa'）
            init_lora_weights = getattr(config, 'init_lora_weights', None)

            # 如果 init_lora_weights 是字符串（说明启用了某种特定 LoRA 方法）
            if isinstance(init_lora_weights, str):

                # 拷贝一份配置，防止后续修改 config 引起副作用
                config = copy(config)

                # 创建保存路径 output_dir/converted，如果已存在则不报错
                os.makedirs(os.path.join(output_dir, 'converted'), exist_ok=True)

                # 如果使用了 LoRA-GA 方法（需要特殊上下文支持）
                if 'lora-ga' in init_lora_weights:
                    try:
                        # 尝试引入 LoRA-GA 的上下文管理器（该库并非 PEFT 官方支持）
                        from lora_ga.entrypoint import LoraGAContext

                        # 使用 LoraGAContext 上下文封装模型，进行 LoRA-GA 权重融合转换
                        with LoraGAContext(model):
                            model.save_pretrained(
                                os.path.join(output_dir, 'converted', 'default'),  # 保存路径为 converted/default
                                path_initial_model_for_weight_conversion=os.path.join(
                                    os.path.dirname(output_dir), 'initial_model'),  # 提供初始模型路径作为基线
                            )
                            # 恢复配置（因为可能在上下文或保存中被修改）
                            model.peft_config['default'] = config

                    except ImportError as e:
                        # 如果未安装 LoRA-GA 库，提示用户安装 GitHub 项目
                        error_message = """
                        Since 'LoRA-GA' is not implemented by PEFT, you will need to install it directly from GitHub.
                        Command: 'pip install git+https://github.com/lxline/LoRA-GA.git'.
                        """
                        logger.info(error_message)
                        # 抛出异常，阻止后续执行
                        raise RuntimeError(error_message) from e

                # 如果是 pissa 或 olora 方法，不需要额外上下文，直接保存转换后的模型
                elif 'pissa' in init_lora_weights or 'olora' in init_lora_weights:
                    model.save_pretrained(
                        os.path.join(output_dir, 'converted', 'default'),  # 保存路径
                        path_initial_model_for_weight_conversion=os.path.join(
                            os.path.dirname(output_dir), 'initial_model'),  # 初始权重路径
                    )
                    # 恢复配置，避免其他逻辑读取到被改动过的 config
                    model.peft_config['default'] = config

    def _load_optimizer_and_scheduler(self, *args, **kwargs):
        """
        加载优化器和学习率调度器。

        功能说明：
        - 调用父类方法以加载训练快照中保存的优化器与调度器；
        - 如果启用了多进程混合精度 (MP) + DDP 分布式（如 DeepSpeed Zero3），
          则对 AdamW 优化器中的 step 状态项做设备修复，防止设备不一致导致训练报错。
        """

        # 调用父类的优化器/调度器加载方法（通常从 checkpoint 恢复）
        super()._load_optimizer_and_scheduler(*args, **kwargs)

        # 如果当前训练启用了多进程 + 分布式（例如使用 DeepSpeed Zero3 或 Megatron）
        if is_mp_ddp():
            # 修复在 MP + DDP 模式下，AdamW 优化器中的状态张量可能跨多个设备的问题
            for v in self.optimizer.state.values():
                # 检查当前参数状态中是否包含 'step'（表示优化器步数）
                if 'step' in v:
                    # 检查除了 step 本身所在设备（和 CPU）之外，还有没有其他设备存在
                    device_set = set([t.device for t in v.values()]) - {v['step'].device, torch.device('cpu')}

                    # 如果发现存在跨设备状态（即 device_set 非空），可能会导致报错
                    if len(device_set) >= 1:
                        # 强制将 step 张量移到 CPU，确保状态字典统一，防止设备不一致错误
                        v['step'] = v['step'].to('cpu')

    def _save_model(self, output_dir: Optional[str] = None, state_dict=None):
        """
        保存模型权重的核心实现。

        功能说明：
        - 根据模型类型，选择对应保存方式；
        - 支持 HuggingFace、LoRA (PEFT)、Reward Model 等；
        - 支持 safetensors 格式；
        - 对 SentenceTransformer 等特殊模型进行定制保存；
        - 支持微软 MS 模型和其他自定义 Adapter 保存。
        """

        # 定义支持 save_pretrained 的模型类型
        supported_classes = (SwiftModel, PreTrainedModel, PeftModel)
        supported_names = ('SentenceTransformer')

        # 如果使用的是 Value Head（如 PPO 模型），加入其类支持列表
        if AutoModelForCausalLMWithValueHead is not None:
            supported_classes = supported_classes + (AutoModelForCausalLMWithValueHead,)

        # 是否使用 safetensors 格式保存
        save_safetensors = self.args.save_safetensors

        # 如果当前模型类型不是标准支持类，也不是特殊名字（如 SentenceTransformer）
        if not isinstance(self.model, supported_classes) and self.model.__class__.__name__ not in supported_names:
            # 如果没有传入 state_dict，则自动调用 model.state_dict()
            if state_dict is None:
                state_dict = self.model.state_dict()

            # 尝试解开封装后的模型再判断是否支持 save_pretrained
            _unwrap_model = unwrap_model(self.model)
            if isinstance(_unwrap_model, supported_classes):
                # 如果是支持类型，使用 save_pretrained 保存权重
                _unwrap_model.save_pretrained(output_dir, state_dict=state_dict, safe_serialization=save_safetensors)
            else:
                # 否则只能保存 state_dict 文件
                logger.info('Trainer.model is not a `PreTrainedModel`, only saving its state dict.')
                if save_safetensors:
                    safetensors.torch.save_file(state_dict, os.path.join(output_dir, 'model.safetensors'))
                else:
                    torch.save(state_dict, os.path.join(output_dir, 'pytorch_model.bin'))

        # 如果是 reward model（如 PPO 训练的 ValueHead 模型）
        elif AutoModelForCausalLMWithValueHead and isinstance(self.model, AutoModelForCausalLMWithValueHead):
            state_dict = self.model.state_dict()
            decoder_state_dict, v_head_state_dict = {}, {}

            # 将 value head 与原模型部分拆分开
            for name, param in state_dict.items():
                if name.startswith('v_head.'):
                    v_head_state_dict[name] = param
                else:
                    decoder_state_dict[name.replace('pretrained_model.', '', 1)] = param

            # 保存 base model
            self.model.pretrained_model.save_pretrained(
                output_dir,
                state_dict=decoder_state_dict or None,
                safe_serialization=save_safetensors
            )

            # 保存 value head 参数
            if save_safetensors:
                from safetensors.torch import save_file
                save_file(
                    v_head_state_dict, os.path.join(output_dir, 'value_head.safetensors'),
                    metadata={'format': 'pt'})
            else:
                torch.save(v_head_state_dict, os.path.join(output_dir, 'value_head.bin'))

        # 微软专属模型的保存方式（例如 InternLM、ChatGLM 等）
        elif is_instance_of_ms_model(self.model):
            PreTrainedModel.save_pretrained(
                self.model,
                output_dir,
                state_dict=state_dict,
                safe_serialization=save_safetensors
            )

        # 使用了自定义 Tuner（如 LoRA、Q-LoRA、AdaLoRA 等）
        elif self.args.train_type in extra_tuners:
            extra_tuners[self.args.train_type].save_pretrained(
                self.model,
                output_dir,
                state_dict=state_dict,
                safe_serialization=save_safetensors
            )

        # 标准 HuggingFace 或 PeftModel 保存
        else:
            if self.model.__class__.__name__ != 'SentenceTransformer':
                self.model.save_pretrained(
                    output_dir,
                    state_dict=state_dict,
                    safe_serialization=save_safetensors
                )
            else:
                # SentenceTransformer 保存方式特殊，需处理其 auto_model 的 state_dict
                @contextmanager
                def save_context():
                    save_pretrained = self.model[0].auto_model.save_pretrained
                    _state_dict = {
                        key[len('0.auto_model.'):] if 'auto_model' in key else key: value
                        for key, value in state_dict.items()
                    }

                    # 修改 save_pretrained 方法，使其接受 state_dict 参数
                    self.model[0].auto_model.save_pretrained = partial(
                        self.model[0].auto_model.save_pretrained,
                        state_dict=_state_dict
                    )
                    yield
                    self.model[0].auto_model.save_pretrained = save_pretrained

                # 执行保存上下文
                with save_context():
                    self.model.save_pretrained(output_dir, safe_serialization=save_safetensors)

                    # 同时将 .py 和 .json 文件也从模型目录复制过去
                    from swift.utils import copy_files_by_pattern
                    copy_files_by_pattern(self.model.model_dir, output_dir, '*.py')
                    copy_files_by_pattern(self.model.model_dir, output_dir, '*.json')

    def _save(self, output_dir: Optional[str] = None, state_dict=None):
        """
        保存整个训练快照，包括模型、训练参数、转换权重、配置和预测结果。

        功能说明：
        - 保存模型权重（通过 _save_model）；
        - 保存训练参数（training_args.bin）；
        - 保存转换后的 LoRA 等模型权重；
        - 复制 args.json；
        - 移动 predict.jsonl；
        - 对非 Adapter 模型（非 LoRA/PEFT）保存 tokenizer、配置等额外文件。
        """

        """Compatible with swift and peft"""
        # 本函数只在主进程执行，不做多进程判断（默认是 rank 0）

        # 如果没有显式指定 output_dir，则使用训练参数中的默认路径
        output_dir = output_dir if output_dir is not None else self.args.output_dir
        os.makedirs(output_dir, exist_ok=True)  # 创建保存目录

        # 1. 保存模型权重（支持多种模型类型）
        self._save_model(output_dir, state_dict)

        # 2. 保存训练参数为二进制文件，便于恢复训练
        torch.save(self.args, os.path.join(output_dir, 'training_args.bin'))

        # 3. 如果使用了 LoRA/pissa/olora 等结构，保存转换后的模型权重
        self._save_converted_model(output_dir)

        # 4. 尝试复制父目录下的 args.json 到输出目录，保留训练参数记录
        args_path = os.path.join(os.path.dirname(output_dir), 'args.json')
        if os.path.exists(args_path):
            shutil.copy(args_path, os.path.join(output_dir, 'args.json'))

        # 5. 移动父目录下的预测文件（推理输出）到输出目录下，便于归档
        predict_jsonl = os.path.join(os.path.dirname(output_dir), 'predict.jsonl')
        if os.path.exists(predict_jsonl):
            shutil.move(predict_jsonl, os.path.join(output_dir, 'predict.jsonl'))

        # 判断是否是 Adapter 模型（SwiftModel 或 PeftModel）
        is_adapter = isinstance(self.model, (SwiftModel, PeftModel))

        # 6. 如果不是 adapter 模型，还需要保存 tokenizer、模型文件等额外资源
        if not is_adapter:
            from swift.llm import save_checkpoint

            # 获取需要额外保存的文件（如 config.json、tokenizer.json 等）
            additional_saved_files = self.model_meta.additional_saved_files

            # 调用 Swift 的通用 checkpoint 保存函数，支持复制 tokenizer、预处理器等
            save_checkpoint(
                None,  # 不保存训练状态（如优化器等）
                self.template.processor,  # 包含 tokenizer、prompt processor 等
                output_dir,
                model_dirs=[self.model.model_dir],
                additional_saved_files=additional_saved_files
            )

            # 7. 如果有原始的 generation_config（文本生成配置），也保存到模型目录
            if getattr(self.model, 'origin_generation_config', None):
                self.model.origin_generation_config.save_pretrained(output_dir)

    def _fix_zero3_gather_all_parameters(self) -> None:
        """
        修正 DeepSpeed Zeor3 模式下获取所有参数时的异常行为。

        背景：
        - DeepSpeed Zero Stage 3 会将模型参数分散在多个 GPU 上；
        - 聚合参数时，有些冻结参数（如 Adapter、LoRA）可能导致不一致或异常；
        - 本方法通过 monkey patch 替换内部聚合逻辑，强制处理 exclude_frozen_parameters 标志。

        功能：
        - 替换 DeepSpeed 的 `_zero3_consolidated_16bit_state_dict` 方法；
        - 避免由于冻结参数未参与训练而造成的聚合错误；
        - 支持 SwiftModel 和 PeftModel 的 LoRA 等结构。
        """

        # 如果启用了 DeepSpeed Zero3，且还没有做过 monkey patch 替换
        if is_deepspeed_zero3_enabled() and not hasattr(self.deepspeed, '_zero3_consolidated_16bit_state_dict_origin'):

            # 获取该方法的函数签名，确认是否支持 exclude_frozen_parameters 参数
            parameters = inspect.signature(self.deepspeed._zero3_consolidated_16bit_state_dict).parameters
            if 'exclude_frozen_parameters' in parameters:

                # 定义新的 monkey patched 方法
                def _zero3_consolidated_16bit_state_dict(model, exclude_frozen_parameters=False):
                    # 解包模型（去除 DDP、混合精度等封装）
                    unwrapped = unwrap_model(model)

                    # 默认排除冻结参数（如 Adapter、LoRA）
                    exclude_frozen_parameters = False
                    if isinstance(unwrapped, SwiftModel) and unwrapped.has_additional_modules:
                        exclude_frozen_parameters = True
                    if isinstance(unwrapped, PeftModel):
                        exclude_frozen_parameters = True

                    # 调用原始 DeepSpeed 方法，传入处理后的参数
                    return model._zero3_consolidated_16bit_state_dict_origin(exclude_frozen_parameters)

                # 先保存原始的 DeepSpeed 内部方法（便于调用）
                self.deepspeed._zero3_consolidated_16bit_state_dict_origin = (
                    self.deepspeed._zero3_consolidated_16bit_state_dict)

                # 替换成 monkey patched 新方法
                self.deepspeed._zero3_consolidated_16bit_state_dict = MethodType(
                    _zero3_consolidated_16bit_state_dict, self.deepspeed)

    def _save_checkpoint(self, *args, **kwargs):
        """
        保存模型检查点（checkpoint）。

        功能说明：
        - 用于在训练过程中定期保存模型、优化器、调度器状态等；
        - 特别处理 DeepSpeed Zero3 模式下的参数聚合；
        - 更新当前模型的检查点路径，供 resume 或后续加载使用；
        - 支持继承自 HuggingFace Trainer 的 checkpoint 保存逻辑。
        """

        # 设置当前步骤对应的 checkpoint 路径，例如：output_dir/checkpoint-1000
        self.state.last_model_checkpoint = os.path.join(
            self.args.output_dir, f'checkpoint-{self.state.global_step}'
        )

        # 修复 DeepSpeed Zero Stage 3 模式下聚合所有参数的问题（冻结参数等）
        self._fix_zero3_gather_all_parameters()

        # 调用父类的 _save_checkpoint 方法，进行实际保存操作（包括模型、优化器、调度器等）
        result = super()._save_checkpoint(*args, **kwargs)

        # 记录日志，提示保存了检查点
        logger.info(f'Saving model checkpoint to {self.state.last_model_checkpoint}')

        # 返回保存结果（通常为 checkpoint 路径或状态信息）
        return result

    @staticmethod
    @contextmanager
    def _fix_grad_norm_nan():
        """
        上下文管理器，修复梯度范数计算出现 NaN 时的异常。

        - 替换 accelerate 中的 clip_grad_norm_ 方法，若出现 NaN，
          则清空所有参数的梯度，防止更新。
        - 退出上下文时恢复原函数。
        """
        from accelerate import Accelerator
        # 保存 accelerate 库中原始的 clip_grad_norm_ 方法
        origin_clip_grad_norm_ = Accelerator.clip_grad_norm_

        # 定义新的 clip_grad_norm_ 函数，用于替换原函数
        def clip_grad_norm_(self, parameters, *args, **kwargs):
            # 将参数转换成列表，方便遍历
            parameters = list(parameters)
            # 调用原始 clip_grad_norm_ 计算梯度范数
            grad_norm = origin_clip_grad_norm_(self, parameters, *args, **kwargs)
            # 如果梯度范数是张量且为 NaN（异常情况）
            if isinstance(grad_norm, torch.Tensor) and grad_norm.isnan().item():
                # 遍历所有参数，将其梯度清空（防止异常更新）
                for p in parameters:
                    p.grad = None
            # 返回计算得到的梯度范数
            return grad_norm

        # 将 accelerate 的 clip_grad_norm_ 替换为自定义的修复版函数
        Accelerator.clip_grad_norm_ = clip_grad_norm_
        try:
            # 进入上下文时，执行用户代码
            yield
        finally:
            # 退出上下文时，恢复 accelerate 中原始的 clip_grad_norm_ 函数
            Accelerator.clip_grad_norm_ = origin_clip_grad_norm_

    def train(self, *args, **kwargs):
        """
        扩展训练方法。

        - 如果是多模态模型，注册后置编码钩子（post_encode hook）。
        - 保存初始模型。
        - 使用 hub 补丁和梯度 NaN 保护上下文执行父类 train。
        - 训练结束后移除钩子。
        - 返回训练结果。
        """
        # 如果模型是多模态的，找到所有相关子模型（model, ref_model, reward_model, value_model）
        if self.model_meta.is_multimodal:
            models = list(
                set([
                    v for k, v in self.__dict__.items()
                    if isinstance(v, nn.Module) and k in {'model', 'ref_model', 'reward_model', 'value_model'}
                ])
            )
            # 注册后置编码钩子，方便多模态模型在编码后进行处理
            self.template.register_post_encode_hook(models)
            logger.info(f'Successfully registered post_encode hook: {[model.__class__.__name__ for model in models]}.')

        # 保存训练开始前的初始模型权重
        self._save_initial_model(self.args.output_dir)

        # 在 hub 补丁环境和梯度范数 NaN 保护上下文中执行父类的训练方法
        with self.hub.patch_hub(), self._fix_grad_norm_nan():
            res = super().train(*args, **kwargs)

        # 训练结束后，移除之前注册的后置编码钩子
        self.template.remove_post_encode_hook()

        # 返回训练结果
        return res

    def push_to_hub(self, *args, **kwargs):
        """
        将模型推送到模型库。

        - 使用 hub 补丁上下文，调用父类推送。
        """
        # 在 hub 的补丁上下文中执行，确保推送时支持自定义逻辑或修正
        with self.hub.patch_hub():
            # 调用父类的 push_to_hub 方法，完成模型推送
            return super().push_to_hub(*args, **kwargs)

    def get_max_cuda_memory(self, device: Optional[Union[torch.device, int]] = None) -> float:
        """
        获取当前训练过程中显存的最大使用量（GiB）。

        - 若 device 未指定，统计所有 GPU 的最大显存占用。
        - 计算所有设备最大显存占用总和，转换为 GiB 单位。
        - 更新并保存 self.max_memory 为历史最高显存。
        - 返回当前最大显存（GiB）。
        """
        if device is None:
            # 未指定设备时，遍历所有 GPU 设备统计最大显存
            mems = [torch.cuda.max_memory_reserved(device=device) for device in range(get_device_count())]
        else:
            # 指定设备时，只获取该设备的最大显存
            mems = [torch.cuda.max_memory_reserved(device=device)]
        # 计算所有设备最大显存总和，单位转换为 GiB
        mem = sum(mems) / 1024 ** 3
        # 更新历史最大显存值
        self.max_memory = max(self.max_memory, mem)
        # 返回当前最大显存（GiB）
        return mem

    def _maybe_log_save_evaluate(self, tr_loss, *args, **kwargs):
        """
        训练过程中，按日志策略打印日志、保存模型、评估。

        - 计算平均训练损失。
        - 统计并重置自定义指标。
        - 记录梯度范数和当前学习率。
        - 记录显存使用和训练速度。
        - 若配置启用 evalscope，则调用专门的评估逻辑。
        - 调用父类对应方法继续处理日志、保存和评估。
        """
        if self.control.should_log and self.state.global_step > self._globalstep_last_logged:
            # 防止重复日志
            self.control.should_log = False

            # all_gather + mean() 跨进程汇总训练损失，取平均值
            tr_loss_scalar = self._nested_gather(tr_loss).mean().item()
            # 计算这段时间内的平均 loss
            loss = tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged)
            logs: Dict[str, float] = {'loss': loss}  # 先记录 loss

            # 统计所有自定义指标，并格式化加入日志
            for k, metric in self._custom_metrics.items():
                value = metric.compute()
                if len(value) == 1:
                    val = list(value.values())[0]
                    logs[k] = val
                else:
                    for k_suffix, val in value.items():
                        new_k = f'{k}_{k_suffix}'
                        logs[new_k] = val
                metric.reset()  # 重置指标，准备下次统计

            # 仅 transformers 4.38+ 版本支持传入 grad_norm 参数
            if version.parse(transformers.__version__) >= version.parse('4.38'):
                grad_norm = args[0]
                if grad_norm is not None:
                    logs['grad_norm'] = grad_norm.item() if isinstance(grad_norm, torch.Tensor) else grad_norm
            # 记录当前学习率
            logs['learning_rate'] = self._get_learning_rate()

            # 如果不是在华为 NPU 设备上，记录 GPU 显存使用情况（单位 GiB）
            if not is_torch_npu_available():
                logs['memory(GiB)'] = round(self.get_max_cuda_memory(), 2)

            # 计算训练速度（步数/秒）
            elapse_time = time.time() - self.start_time
            logs['train_speed(iter/s)'] = round(self.state.global_step / elapse_time, 6)

            # 删除日志中值为 None 的条目
            for k in list(logs.keys()):
                if logs[k] is None:
                    logs.pop(k)

            # 重置 tr_loss（防止累加）
            tr_loss -= tr_loss
            # 累加总损失
            self._total_loss_scalar += tr_loss_scalar
            # 更新日志时间戳
            self._globalstep_last_logged = self.state.global_step
            # 记录 FLOPS
            self.store_flos()
            # 打印日志
            self.log(logs)

        # 如果配置启用了 evalscope，且需要评估，则调用对应评估函数
        if self.args.eval_use_evalscope and self.control.should_evaluate:
            self._evalscope_eval()

        # 调用父类的方法继续日志保存和评估
        super()._maybe_log_save_evaluate(tr_loss, *args, **kwargs)

    def create_optimizer_and_scheduler(self, num_training_steps: int):
        """
        创建优化器和学习率调度器。

        - 支持自定义优化器映射，通过 swift.plugin.optimizers_map 获取对应优化器创建函数。
        - 若自定义优化器创建函数返回 None，调用默认优化器和调度器创建逻辑。
        """
        if self.args.optimizer is not None:
            # 从 swift.plugin 中获取自定义优化器映射表
            from swift.plugin import optimizers_map
            # 获取指定优化器的回调函数
            optimizer_callback = optimizers_map[self.args.optimizer]
            # 通过回调函数创建优化器和学习率调度器
            self.optimizer, self.lr_scheduler = optimizer_callback(self.args, self.model, self.train_dataset)
            # 如果优化器未创建，则调用默认创建函数
            if self.optimizer is None:
                self.create_optimizer()
            # 如果学习率调度器未创建，则调用默认创建函数
            if self.lr_scheduler is None:
                self.create_scheduler(num_training_steps=num_training_steps, optimizer=self.optimizer)
        else:
            # 若未指定自定义优化器，调用父类默认逻辑
            super().create_optimizer_and_scheduler(num_training_steps=num_training_steps)

    def _compute_acc(self, outputs, labels) -> None:
        """
        计算分类任务的准确率。

        - 每 acc_steps 步计算一次准确率。
        - 使用 torchacc 进行加速计算（如果可用）。
        - 计算指标并存储到自定义指标容器中。
        """
        args = self.args
        acc_steps = args.acc_steps  # 获取累积步数，用于决定何时计算准确率
        preds = outputs.logits.argmax(dim=-1)  # 预测结果取最大概率的类别索引
        # 仅在当前训练步数是 acc_steps 的整数倍时计算准确率
        if self.state.global_step % acc_steps == 0:
            # 如果可用，调用 torchacc 进行图修剪等加速操作
            if use_torchacc():
                ta_trim_graph()  # 裁剪计算图
                preds = preds.to('cpu')  # 将预测结果移至CPU，方便后续处理
                labels = labels.to('cpu')  # 将标签移至CPU
            # 调用 compute_acc 函数计算准确率指标，传入当前配置参数
            metrics = compute_acc(
                preds, labels, acc_strategy=args.acc_strategy, is_encoder_decoder=self.template.is_encoder_decoder)
            # 将计算得到的各项指标更新到自定义指标集合中，方便后续日志输出和分析
            for k, v in metrics.items():
                if k not in self._custom_metrics:
                    self._custom_metrics[k] = MeanMetric(nan_value=None)  # 创建指标容器
                self._custom_metrics[k].update(v)  # 更新指标值

    @torch.no_grad()
    def _evalscope_eval(self):
        """
        基于 evalscope 框架的自定义评估。

        - 创建 EvalModel，配置评估任务。
        - 执行评估并收集分数。
        - 记录日志。
        - 切回训练模式。
        - 返回评估结果字典。
        """
        from ..llm.eval.utils import EvalModel
        from evalscope import TaskConfig, run_task
        from evalscope.constants import EvalType

        self.model.eval()  # 切换模型到评估模式，关闭 dropout 等训练专用功能
        max_batch_size = self.args.per_device_eval_batch_size  # 评估时的最大批量大小

        # 创建自定义 EvalModel，用于兼容 evalscope 评估流程
        custom_model = EvalModel(
            self.model, self.template, max_batch_size=max_batch_size, model_name=f'model-step{self.state.global_step}')

        # 配置评估任务参数，包括数据集、评估类型、工作目录等
        task_config = TaskConfig(
            model=custom_model,
            eval_type=EvalType.CUSTOM,
            datasets=self.args.eval_datasets,
            dataset_args=self.args.eval_datasets_args,
            limit=self.args.eval_limit,
            work_dir=os.path.join(self.args.output_dir, 'eval'),
            eval_batch_size=max_batch_size,
            generation_config=self.args.eval_generation_config or {'max_tokens': 512},
        )

        # 开始执行评估任务，返回评估报告
        eval_report = run_task(task_config)

        # 将评估报告转为字典格式，键加前缀 'test_'，值为分数
        eval_dict = {f'test_{k}': v.score for k, v in eval_report.items()}

        self.log(eval_dict)  # 记录日志

        self.model.train()  # 恢复训练模式

        return eval_dict  # 返回评估结果字典

    def get_batch_samples(self, *args, **kwargs):
        """
        获取训练批次样本。

        - 调用父类方法获取样本。
        - 如果启用 sequence_parallel_size > 1，执行分布式 all_reduce 合并样本数量。
        """
        # 调用父类获取批次样本及其数量
        res = super().get_batch_samples(*args, **kwargs)
        # 如果没有启用序列并行，直接返回结果
        if self.template.sequence_parallel_size == 1:
            return res
        # 否则，拆包批次样本和批次数量
        batch_samples, num_items_in_batch = res
        # 分布式环境下对批次数量进行 all_reduce 求和，保证全局批次数一致
        dist.all_reduce(num_items_in_batch, dist.ReduceOp.SUM)
        # 返回合并后的批次样本和批次数量
        return batch_samples, num_items_in_batch


class DataLoaderMixin:
    """
    DataLoader 混入类，提供训练和评估的数据加载器获取方法。

    训练数据加载器（get_train_dataloader）支持：
    - 多模态序列并行（sequence_parallel_size > 1）时，使用专用的 sequence_parallel dataloader。
    - 否则，根据训练数据集类型（可测长度或 IterableDataset），
      采用 BatchSamplerShard + DataLoaderShard 或普通 DataLoader + DataLoaderDispatcher。
    - 配置 dataloader 参数（collate_fn、num_workers、pin_memory 等），
      并根据是否分布式环境调整预取因子。

    评估数据加载器（get_eval_dataloader）支持：
    - 多模态序列并行时，使用 sequence_parallel 的 dataloader。
    - 否则调用父类实现。
    """

    def get_train_dataloader(self):
        """
        获取训练数据加载器。

        - 若启用了序列并行且大小 > 1，使用专门的 sequence_parallel dataloader。
        - 否则根据训练数据集类型，选择不同的加载器：
          - 可测长度的数据集，使用 BatchSamplerShard + DataLoaderShard，提高数据加载效率。
          - IterableDataset，调整预取因子并使用 DataLoader + DataLoaderDispatcher 进行封装。
        - 配置 dataloader 的参数，如 collate_fn、num_workers、pin_memory 等。
        - 在分布式环境下，自动调整预取因子以提高数据加载性能。
        - 如果未提供训练数据集，则抛出异常。
        """
        dataloader = None
        # 如果启用序列并行且大小 > 1，使用专门的 sequence_parallel dataloader
        if self.template.sequence_parallel_size > 1:
            from swift.trainers.sequence_parallel import sequence_parallel
            dataloader = sequence_parallel.get_dataloader(self, self.train_dataset, self._train_batch_size)

        # 如果未创建 dataloader，则创建普通 dataloader
        if dataloader is None:
            if self.train_dataset is None:
                raise ValueError('Trainer: training requires a train_dataset.')
            args = self.args
            train_dataset = self.train_dataset

            # 配置 dataloader 参数
            dataloader_params = {
                'collate_fn': self.data_collator,
                'num_workers': args.dataloader_num_workers,
                'pin_memory': args.dataloader_pin_memory,
                'persistent_workers': args.dataloader_persistent_workers,
                'prefetch_factor': args.dataloader_prefetch_factor
            }
            batch_sampler_params = {
                'drop_last': args.dataloader_drop_last,
                'shuffle': args.train_dataloader_shuffle,
                'data_seed': args.data_seed,
            }

            # 如果数据集可测长度，使用 BatchSamplerShard + DataLoaderShard
            if hasattr(train_dataset, '__len__'):
                batch_sampler = BatchSamplerShard(
                    len(train_dataset), batch_size=self._train_batch_size, **batch_sampler_params)
                dataloader = DataLoaderShard(train_dataset, batch_sampler, **dataloader_params)
            else:
                # 对 IterableDataset，调整 prefetch_factor 并用 DataLoaderDispatcher 包装
                if dist.is_initialized():
                    dataloader_params['prefetch_factor'] = dataloader_params['prefetch_factor'] * dist.get_world_size()
                dataloader = DataLoader(train_dataset, batch_size=self._train_batch_size, **dataloader_params)
                dataloader = DataLoaderDispatcher(dataloader)

        return dataloader

    def get_eval_dataloader(self, eval_dataset=None):
        """
        获取评估数据加载器。

        - 如果启用了序列并行且大小 > 1，则使用专门的 sequence_parallel dataloader。
        - 如果未指定 eval_dataset，且实例属性 eval_dataset 也为空，则抛出异常。
        - 否则调用父类的 get_eval_dataloader 方法获取评估数据加载器。
        """
        dataloader = None
        # 如果启用序列并行，使用 sequence_parallel dataloader
        if self.template.sequence_parallel_size > 1:
            from swift.trainers.sequence_parallel import sequence_parallel
            if eval_dataset is None and self.eval_dataset is None:
                raise ValueError('Trainer: evaluation requires an eval_dataset.')
            eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
            dataloader = sequence_parallel.get_dataloader(self, eval_dataset, self.args.eval_batch_size)

        # 如果未创建 dataloader，则调用父类方法
        if dataloader is None:
            return super().get_eval_dataloader(eval_dataset=eval_dataset)
        return dataloader
